/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;
import java.util.concurrent.atomic.AtomicLong;

import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.util.*;
import edu.umd.clip.lm.util.tree.BinaryTree;
import edu.umd.clip.lm.model.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class TrellisSlice {
	//private List<SliceItem> elements; 
	protected HashMap<SliceItem,ArrayList<SliceItem>> curr2prev;
	protected HashMap<SliceItem,ArrayList<SliceItem>> prev2curr;
	protected HashMap<ClusterId, SliceItem> itemsByCluster;
	protected double bestProb = 0.0;
	private long overtFactors;
	private BinaryPrefix wordPrefix;
	private Trellis trellis;
	private int time;
	private double totalProb = Double.NaN;
	private boolean dummy = false; // whether the slice is just used to prefix p(w t | ctx) 
	private FuturePredictions predictions;
	private double sliceLogScale;
	private double cumulativeLogScale;
	
	/**
	 * @param overtFactors
	 * @param trellis
	 * @param time
	 */
	public TrellisSlice(long overtFactors, Trellis trellis, int time) {
		this.overtFactors = overtFactors;
		Experiment experiment = Experiment.getInstance();
		int value = FactorTuple.getValue(overtFactors, experiment.getTupleDescription().getMainFactorIndex());
		wordPrefix = experiment.getWordPrefix(value);

		this.trellis = trellis;
		this.time = time;
		//this.elements = new LinkedList<SliceItem>();
		this.prev2curr = new HashMap<SliceItem,ArrayList<SliceItem>>(100);
		this.curr2prev = new HashMap<SliceItem,ArrayList<SliceItem>>(100);
		this.itemsByCluster = new HashMap<ClusterId, SliceItem>(200);
		predictions = new FuturePredictions();
	}
	
	static protected class SliceItem implements AbstractProbTree, Cloneable {
		protected ProbTree tree;
		private double prevProb;
		protected ClusterId clusterid;
		protected TrellisSlice slice;
		private static AtomicLong lastid = new AtomicLong();
		private long id;
		//private LinkQueue queue;
		protected BinaryTree<HistoryTreePayload> decodingPosition[];
		/**
		 * @param tree
		 * @param prevProb
		 * @param clusterid
		 */
		public SliceItem(ProbTree tree, double prevProb, ClusterId clusterid, BinaryTree<HistoryTreePayload>[] positions) {
			this.tree = tree;
			this.prevProb = prevProb;
			this.clusterid = clusterid;
			this.decodingPosition = positions;
			id = makeId();
		}

		@SuppressWarnings("unchecked")
		public SliceItem(ProbTree tree, double prevProb, ClusterId clusterid) {
			this(tree, prevProb, clusterid, new BinaryTree[Experiment.getInstance().getNumLMs()]);
		}
		
		public SliceItem(ProbTree tree, double prevProb) {
			this(tree, prevProb, new ClusterId());
		}		
		
		public BinaryTree<HistoryTreePayload> getPosition(int lmId) {
			return decodingPosition[lmId];
		}
		
		public void setPosition(int lmId, BinaryTree<HistoryTreePayload> position) {
			decodingPosition[lmId] = position;
		}
		
		public void addPrevProb(double prob) {
			prevProb += prob;
		}
		
		public double getProb() {
			return prevProb * tree.getTotalProb();
		}
		
		public double getPrevProb() {
			return prevProb;
		}
		
		public void setPrevProb(double p) {
			prevProb = p;
		}
		
		public void updatePrevProb() {
			ArrayList<SliceItem> prevs = slice.curr2prev.get(this);
			if (prevs != null) {
				double prevProb = 0;
				for(SliceItem item : prevs) {
					prevProb += item.prevProb * item.tree.getProbTree().getTotalProb();
				}
				this.prevProb = prevProb;
			}
		}
		
		private static long makeId() {
			return lastid.incrementAndGet();
		}
		
		public String toString() {
			return Long.toString(getId()) + "at" + Integer.toString(slice.getTime());
		}
		
		public long getId() {
			return id;
			//return 0;
		}

		public BinaryPrefix getWordPrefix() {
			return slice.wordPrefix;
		}

		public Boolean testCut(BinaryPrefix cuttingPrefix) {
			return tree.testCut(cuttingPrefix);
		}

		public long getTuple() {
			return slice.overtFactors;
		}

		protected boolean checkPrevProb() {
			return slice.checkPrevProb(this);
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.AbstractProbTree#testOvertFactors(edu.umd.clip.lm.factors.Factors)
		 */
		public Boolean testOvertFactors(long factors) {
			throw new UnsupportedOperationException();
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.AbstractProbTree#cut(edu.umd.clip.lm.util.BinaryPrefix)
		 */
		public ProbTree cut(BinaryPrefix cuttingPrefix) {
			throw new UnsupportedOperationException();
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.AbstractProbTree#getPrefix()
		 */
		public BinaryPrefix getPrefix() {
			throw new UnsupportedOperationException();
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.AbstractProbTree#getTotalProb()
		 */
		public double getTotalProb() {
			throw new UnsupportedOperationException();
		}

		public ProbTree getProbTree() {
			return tree.getProbTree();
		}

		@Override
		public Object clone() {
			SliceItem copy = new SliceItem((ProbTree)tree.clone(), 0, (ClusterId)clusterid.clone());
			return copy;
		}
		
		public SliceItem copy() {
			SliceItem copy = new SliceItem(tree.getProbTree().copy(), 0, (ClusterId)clusterid.clone());
			return copy;			
		}
		
		protected void makeGhost() {
			// TODO: fix ghosts
			//tree = new GhostProbTree(tree);
			//queue = null;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.AbstractProbTree#scale(double)
		 */
		public void scale(double scale) {
			this.tree.scale(scale);
		}

		/*
		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result
					+ ((clusterid == null) ? 0 : clusterid.hashCode());
			result = prime * result + ((tree == null) ? 0 : tree.hashCode());
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj)
				return true;
			if (obj == null)
				return false;
			if (!(obj instanceof SliceItem))
				return false;
			SliceItem other = (SliceItem) obj;
			if (clusterid == null) {
				if (other.clusterid != null)
					return false;
			} else if (!clusterid.equals(other.clusterid))
				return false;
			if (tree == null) {
				if (other.tree != null)
					return false;
			} else if (!tree.equals(other.tree))
				return false;
			return true;
		}
		*/
	}
	
	protected void insertItem(SliceItem item) {
		//elements.add(item);
		assert(item.slice == null || item.slice == this);
		item.slice = this;
		//item.tree.setTuple(overtFactors);
		if (trellis.getDebug() > 1) {
			System.out.printf("Slice #%d: insertItem %s, prevProb = %e, treeProb = %e\n", time, 
				item, item.getPrevProb(), item.tree.getTotalProb());
		}
		totalProb = Double.NaN;
	}

	protected double computeTotalProb() {
		double prob = 0.0;
		//System.err.println("totalProb, slice #" + Integer.toString(time));
		for(SliceItem item : getElements()) {
			double elementProb = item.getPrevProb() * item.tree.getTotalProb(); 
			//System.err.println("  " + Double.toString(elementProb));
			prob += elementProb;
		}
		return prob;
	}
	
	protected void changeForwardProbability(SliceItem item, double probChange) {
		TrellisSlice nextSlice = getNext();
		if (nextSlice == null) return;
		Collection<SliceItem> nextItems = nextSlice.prev2curr.get(item);
		if (nextItems != null) {
			TrellisSlice nextNextSlice = nextSlice.getNext();
			for(SliceItem i : nextItems) {
				i.addPrevProb(probChange);
				if (nextNextSlice != null) {
					// TODO: this can get ugly in higher order models
					nextSlice.changeForwardProbability(i, probChange * i.tree.getTotalProb());
				}
			}
		}
	}
	
	@SuppressWarnings("null")
	protected void removeNodes(HashSet<SliceItem> nodes) {
		TrellisSlice prevSlice = getPrevious();
		HashSet<SliceItem> prevItems = null;
		if (prevSlice != null) {
			prevItems = new HashSet<SliceItem>();
		}
		
		TrellisSlice nextSlice = getNext();
		HashSet<SliceItem> nextItems = null;
		HashSet<SliceItem> nextToRemove = null;
		if (nextSlice != null) {
			nextItems = new HashSet<SliceItem>();
			nextToRemove = new HashSet<SliceItem>();
		}
		
		for(SliceItem item : nodes) {
			if (prevSlice != null) {
				List<SliceItem> prevs = curr2prev.get(item);
				if (prevs != null) {
					for(SliceItem i : prevs) {
						if (!prevItems.contains(i)) {
							List<SliceItem> lst = prev2curr.get(i);
							lst.removeAll(nodes);
							prevItems.add(i);
						}
					}
					curr2prev.remove(item);
				}
			}
			
			if (nextSlice != null) {
				List<SliceItem> nexts = nextSlice.prev2curr.get(item);
				if (nexts != null) {
					for(SliceItem i : nexts) {
						if (!nextItems.contains(i)) {
							List<SliceItem> lst = nextSlice.curr2prev.get(i);
							lst.removeAll(nodes);
							if (lst.isEmpty()) {
								nextToRemove.add(i);
							}
							nextItems.add(i);
						}
					}
					nextSlice.prev2curr.remove(item);
				}
			}
		}
		//elements.removeAll(nodes);
		if (nextToRemove != null && nextToRemove.size() > 0) {
			nextSlice.removeNodes(nextToRemove);
		}
	}
	
	public double getTotalProb() {
		if (Double.isNaN(totalProb)) {
			totalProb = computeTotalProb();
		}
		return totalProb;
	}
	/**
	 * @return the time
	 */
	public int getTime() {
		return time;
	}
	
	public TrellisSlice getPrevious() {
		return trellis.getPreviousSlice(this);
	}
	
	public TrellisSlice getNext() {
		return trellis.getNextSlice(this);
	}

	/**
	 * @return the elements
	 */
	public Set<SliceItem> getElements() {
		return curr2prev.keySet();
	}
/*
	protected void setElements(List<SliceItem> elements) {
		this.elements = elements;
	}
*/	
	/**
	 * @return the overtFactors
	 */
	public long getOvertFactors() {
		return overtFactors;
	}

	public ArrayList<SliceItem> getPreviousItems(SliceItem item) {
		return curr2prev.get(item);
	}
	
	protected SliceItem findSameItemInCluster(ClusterId clusterid) {
		SliceItem item = itemsByCluster.get(clusterid);
		return item;
	}
	
	protected void addItemToCluster(ClusterId clusterid, SliceItem item) {
		itemsByCluster.put(clusterid, item);
		if (trellis.getDebug() > 1) {
			System.out.printf("addItemToCluster: item %s, cluster: (%s)\n", 
					item.toString(), clusterid);
		}
	}
	
	protected boolean checkPrevProb(SliceItem item) {
		ArrayList<SliceItem> prevLinks = curr2prev.get(item);
		if (prevLinks == null) return true;
		
		double totalPrev = 0;
		for(SliceItem prev : prevLinks) {
			totalPrev += prev.getPrevProb() * prev.tree.getTotalProb();
		}
		if (!ProbMath.approxEqual(totalPrev, item.getPrevProb())) {
			System.out.printf("%s has wrong prev prob: %e, should be %e\n", item, item.getPrevProb(), totalPrev);
			System.out.print("prev links:");
			for(SliceItem prev : prevLinks) {
				System.out.printf(" %s", prev);
			}
			System.out.println();
			return false;
		}
		return true;
	}
	
	public boolean checkLinks() {
		if (trellis.getDebug() == 0) return true;
		
		TrellisSlice prevSlice = getPrevious();
		for(SliceItem item : getElements()) {
			ArrayList<SliceItem> prevLinks = curr2prev.get(item);
			if (item == null) {
				System.err.println("must have links from previous slice");
				return false;
			}
			for(SliceItem prevItem : prevLinks) {
				if (!prevSlice.getElements().contains(prevItem)) {
					System.err.println("item is not present in the previous slice");
					return false;
				}
				if (!prev2curr.get(prevItem).contains(item)){
					System.err.println("backreference s not present");
					return false;
				}
			}
			
		}
		return true;
	}
	
	public int numItems() {
		return curr2prev.size();
	}

	/**
	 * @return the trellis
	 */
	public Trellis getTrellis() {
		return trellis;
	}

	protected boolean isDummy() {
		return dummy;
	}

	protected void setDummy(boolean dummy) {
		this.dummy = dummy;
	}
	
	protected void makeGhost() {
		for(SliceItem item : getElements()) {
			item.makeGhost();
		}
	}

	public void addPrediction(BinaryTree<HistoryTreePayload>[] nodes, SliceItem prevItem) {
		predictions.addPrediction(nodes, prevItem);
	}

	public FuturePredictions getPredictions() {
		return predictions;
	}
	
	public void clearPredictions() {
		predictions = null;
	}
	
	@SuppressWarnings("unchecked")
	protected void split(SliceItem item, SliceItem cutOff) {
		assert(item.slice == this);
		cutOff.slice = this;
		// copy links from the previous slice
		{
			ArrayList<SliceItem> prevLinks = curr2prev.get(item);
			if (prevLinks != null) {
				curr2prev.put(cutOff, (ArrayList<SliceItem>)prevLinks.clone());
				for(SliceItem i : prevLinks) {
					prev2curr.get(i).add(cutOff);
				}
			}
		}
		// copy links to the following slice if it exists
		{
			TrellisSlice nextSlice = getNext();
			if (nextSlice != null) {
				ArrayList<SliceItem> fwdLinks = nextSlice.prev2curr.get(item);
				if (fwdLinks != null) {
					nextSlice.prev2curr.put(cutOff, (ArrayList<SliceItem>)fwdLinks.clone());
					for(SliceItem i : fwdLinks) {
						nextSlice.curr2prev.get(i).add(cutOff);
					}
				}
			}
		}
		insertItem(cutOff);
	}
	
	protected void updateForwardProbabilities(SliceItem item, double probDiff) {
		assert(item.slice == this);
		TrellisSlice nextSlice = getNext();
		if (nextSlice != null) {
			Collection<SliceItem> prevs = nextSlice.prev2curr.get(item);
			if (prevs != null) {
				for(SliceItem i : prevs) {
					i.addPrevProb(probDiff);
				}
			}
		}
	}

	public double getSliceLogScale() {
		return sliceLogScale;
	}

	public double getCumulativeLogScale() {
		return cumulativeLogScale;
	}
	
	protected void scale(double scale) {
		for(SliceItem item : getElements()) {
			item.scale(scale);
		}
		this.sliceLogScale -= Math.log10(scale);
		
		this.cumulativeLogScale = time > 2 ? getPrevious().cumulativeLogScale : 0.0;
		this.cumulativeLogScale += sliceLogScale;
	}
	
	protected void setCumulativeLogScale(double logScale) {
		cumulativeLogScale = logScale;
	}
	
	protected double getMaxScale() {
		double maxScale = 0.0;
		for(SliceItem item : getElements()) {
			double scale = item.tree.getProbTree().getScale();
			if (scale > maxScale) {
				maxScale = scale;
			}
		}
		return maxScale > 0.0 ? maxScale : 1.0;
	}
}
